package com.zzyang.mybatis.jdbc.core;

import com.zzyang.mybatis.jdbc.spring.SpringContextUtils;
import com.zzyang.mybatis.jdbc.spring.config.DataSourceDefinition;
import com.zzyang.mybatis.jdbc.spring.config.DataSourcePostProcessor;
import com.zzyang.mybatis.jdbc.spring.config.SeparationJdbcProperties;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.util.CollectionUtils;

import javax.sql.DataSource;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * 数据源上下文
 *
 * @author yisuan
 * @date 2021/4/10 08:55
 */
public class DataSourceContext {

    /**
     * 所有数据源id对应对应的数据源
     */
    private Map<String, DataSource> allDataSource = new HashMap<>();

    /**
     * 所有从数据源id对应对应的数据源
     */
    private Map<String, DataSource> slaveDataSource = new HashMap<>();

    /**
     * 可用的从数据源id ,可能会动态改变，目前还没有支持这个功能
     */
    private List<String> availableSlaveDataSourceIds = new CopyOnWriteArrayList<>();

    /**
     * 主数据源id
     */
    private String masterDataSourceId;

    /**
     * 主数据源实例
     */
    private DataSource masterDataSource;

    public DataSourceContext(SeparationJdbcProperties properties) {
        initAllDataSourceMap(properties);
        processDataSource();
        initDataSource(properties.getMasterDataSourceId());
    }

    private void initDataSource(String masterDataSourceId) {
        this.masterDataSourceId = masterDataSourceId;
        for (Map.Entry<String, DataSource> dataSourceEntry : allDataSource.entrySet()) {
            DataSource dataSource = dataSourceEntry.getValue();
            String dataSourceId = dataSourceEntry.getKey();
            if (masterDataSourceId.equals(dataSourceId)) {
                masterDataSource = dataSource;
            } else {
                availableSlaveDataSourceIds.add(dataSourceId);
                slaveDataSource.put(dataSourceId, dataSource);
            }
        }
        if (masterDataSource == null) {
            throw new IllegalArgumentException("根据 masterDataSourceId=" + masterDataSourceId + " 无法查找到对应的主数据源");
        }
    }


    private void initAllDataSourceMap(SeparationJdbcProperties separationJdbcProperties) {
        List<DataSourceDefinition> dataSourceDefinitions = separationJdbcProperties.getDatasources();
        //优先根据配置的信息来确定数据源，当配置的信息不存在时候，那么就从spring中获取数据源信息
        if (CollectionUtils.isEmpty(dataSourceDefinitions)) {
            Map<String, DataSource> tempMap = SpringContextUtils.getBeansOfType(DataSource.class, false, true);
            for (Map.Entry<String, DataSource> stringDataSourceEntry : tempMap.entrySet()) {
                DataSource dataSource = stringDataSourceEntry.getValue();
                if (dataSource instanceof DefaultRoutingDataSource) {
                    //排除 DefaultRoutingDataSource，因为这个是我们自己创建的，不属于真正的数据源
                    continue;
                }
                allDataSource.put(stringDataSourceEntry.getKey(), dataSource);
            }
        } else {
            for (DataSourceDefinition dataSourceDefinition : dataSourceDefinitions) {
                allDataSource.put(dataSourceDefinition.getId(),
                        DataSourceBuilder.create()
                                .type(dataSourceDefinition.getType())
                                .driverClassName(dataSourceDefinition.getDriverClassName())
                                .url(dataSourceDefinition.getUrl())
                                .username(dataSourceDefinition.getUsername())
                                .password(dataSourceDefinition.getPassword())
                                .build());
            }
        }
        if (CollectionUtils.isEmpty(allDataSource)) {
            throw new IllegalArgumentException("没有可用的数据源");
        }
    }

    private void processDataSource() {
        Map<String, DataSourcePostProcessor> postProcessorMap = SpringContextUtils.getBeansOfType(DataSourcePostProcessor.class);
        if (CollectionUtils.isEmpty(postProcessorMap)) {
            return;
        }
        for (Map.Entry<String, DataSource> stringDataSourceEntry : allDataSource.entrySet()) {
            String id = stringDataSourceEntry.getKey();
            DataSource dataSource = stringDataSourceEntry.getValue();
            for (DataSourcePostProcessor dataSourcePostProcessor : postProcessorMap.values()) {
                dataSourcePostProcessor.process(id, dataSource);
            }
        }
    }

    public String getMasterDataSourceId() {
        return masterDataSourceId;
    }

    public DataSource getMasterDataSource() {
        return masterDataSource;
    }

    public List<String> getAvailableSlaveDataSourceIds() {
        return new ArrayList<>(availableSlaveDataSourceIds);
    }

    public Map<String, DataSource> getAllDataSource() {
        return allDataSource;
    }

    public Map<String, DataSource> getSlaveDataSource() {
        return slaveDataSource;
    }

}
